import { Server } from 'http'
import { WebSocket, WebSocketServer } from 'ws'
import { executionEmitter } from './execution'
import { logger } from './logger'

// WebSocket 连接管理
class WebSocketManager {
    private wss: WebSocketServer
    private clients = new Map<string, Set<WebSocket>>()

    constructor(server: Server) {
        this.wss = new WebSocketServer({ server })
        this.setupWebSocket()
        this.setupEventListeners()
    }

    private setupWebSocket() {
        this.wss.on('connection', (ws, req) => {
            const userId = this.getUserIdFromRequest(req)
            if (!userId) {
                ws.close()
                return
            }

            this.addClient(userId, ws)

            ws.on('close', () => {
                this.removeClient(userId, ws)
            })

            ws.on('error', (error) => {
                logger.error('WebSocket error:', error)
                ws.close()
            })
        })
    }

    private setupEventListeners() {
        // 监听执行进度
        executionEmitter.on('progress', (data) => {
            this.broadcastToUser(data.userId, {
                type: 'progress',
                data
            })
        })

        // 监听执行完成
        executionEmitter.on('completed', (data) => {
            this.broadcastToUser(data.userId, {
                type: 'completed',
                data
            })
        })

        // 监听执行失败
        executionEmitter.on('failed', (data) => {
            this.broadcastToUser(data.userId, {
                type: 'failed',
                data
            })
        })
    }

    private getUserIdFromRequest(req: any): string | null {
        // 从请求中获取用户ID (根据你的认证方式实现)
        return req.headers['x-user-id'] || null
    }

    private addClient(userId: string, ws: WebSocket) {
        if (!this.clients.has(userId)) {
            this.clients.set(userId, new Set())
        }
        this.clients.get(userId)!.add(ws)
    }

    private removeClient(userId: string, ws: WebSocket) {
        const userClients = this.clients.get(userId)
        if (userClients) {
            userClients.delete(ws)
            if (userClients.size === 0) {
                this.clients.delete(userId)
            }
        }
    }

    private broadcastToUser(userId: string, message: any) {
        const userClients = this.clients.get(userId)
        if (userClients) {
            const data = JSON.stringify(message)
            for (const client of userClients) {
                if (client.readyState === WebSocket.OPEN) {
                    client.send(data)
                }
            }
        }
    }
}

export const createWebSocketManager = (server: Server) => {
    return new WebSocketManager(server)
} 